//----------------------------------*-C++-*----------------------------------//
/*!
 * \file   SweepSource.cc
 * \author Jeremy Roberts
 * \date   Aug 24, 2011
 * \brief  SweepSource member definitions.
 * \note   Copyright (C) 2011 Jeremy Roberts. 
 */
//---------------------------------------------------------------------------//
// $Rev:: 141                                           $:Rev of last commit
// $Author:: j.alyn.roberts@gmail.com                   $:Author of last commit
// $Date:: 2011-09-14 19:28:23 +0000 (Wed, 14 Sep 2011) $:Date of last commit
//---------------------------------------------------------------------------//

#include <vector>

#include "SweepSource.hh"

namespace slabtran
{

// Constructor.
template <class Equations, class D>
SweepSource<Equations, D>::SweepSource(SP_state state,
                                       SP_mesh mesh,
                                       SP_angularmesh angularmesh,
                                       SP_materials materials,
                                       SP_momentsindex momentsindex,
                                       SP_moment_to_discrete m_operator) :
  d_state(state),
  d_mesh(mesh),
  d_angularmesh(angularmesh),
  d_materials(materials),
  d_momentsindex(momentsindex),
  d_m_operator(m_operator),
  d_within_group_source(mesh, momentsindex->legendre_order(), 0.0),
  d_in_scatter_source(mesh, momentsindex->legendre_order(), 0.0),
  d_source(mesh, Equations::number_unknowns(), 0.0)
{
  Require(d_state);         // Require that all these objects
  Require(d_mesh);          // are built, i.e. none of the
  Require(d_angularmesh);   // pointers are NULL.
  Require(d_materials);
  Require(d_momentsindex);
  Require(d_m_operator);
}


// Set in-scatter source, i.e. s_g = S_gg*phi_g.
template <class Equations, class D>
void SweepSource<Equations, D>::
build_in_scatter_source(int g)
{

  // Make short reference to within group source
  Moments_Field_t &s = d_in_scatter_source;

  // Grab cell material id's.
  const typename Mesh<D>::Vec_Id &matid = d_mesh->get_cell_id(MATERIAL_ID);

  // Set g_prime bounds for downscatter.
  int g_prime_begin = 0;
  int g_prime_end   = g;
  // Modify for adjoint.
  if (d_angularmesh->adjoint())
  {
      g_prime_begin = g + 1;
      g_prime_end   = d_state->number_groups();
  }

  // Loop through downscatter.
  for (int g_prime = g_prime_begin; g_prime < g_prime_end; g_prime++)
  {

    // Grab the g_prime flux moments.
    const Moments_Field_t &phi_g_prime = d_state->get_phi(g_prime);

    // Loop through cells and unknowns so we can calculate S * v
    // iteration-vector-flux moments at each unknown location.
    for (int cell = 0; cell < d_mesh->number_cells(); cell++)
    {
      // Grab the cross sections for this cell.
      const typename Materials::XS &xs = d_materials->get(matid[cell]);

      // Check if we have downscatter in from g_prime to g.  If
      // not, then we continue to the next cell.  We could do all groups for
      // a cell, but that probably is way less cache-efficient.
      if (xs.group_lower_bound(g) > g_prime) continue;

      // Loop through spatial unknowns in the cell.
      for (int n = 0, N = Equations::number_unknowns(); n < N; n++)
      {
        // Loop through moments.
        for (int l = 0; l <= phi_g_prime.legendre_order(); l++)
        {
          // Get sigma_s for g<-g_prime, l.
          double sig = xs.sigma_s(g, g_prime, l);

          // Get the m indices for this l.
          const std::vector<int> &m_index = d_momentsindex->m_index(l);

          // Loop over m.
          for (int i = 0; i < m_index.size(); i++)
          {
            int m = m_index[i];
            // Add the prime contribution.
            s(cell, n, l, m) += sig * phi_g_prime(cell, n, l, m);
          }
        }

      }
    } // end cell loop
  } // end downscatter loop

  // Return if we choose to skip upscatter.
  if (1==0) return; // \todo add upscatter logic

  // Set g_prime bounds for upscatter
  g_prime_begin = g + 1;
  g_prime_end   = d_state->number_groups();
  // Modify for adjoint
  if (d_angularmesh->adjoint())
  {
      g_prime_begin = 0;
      g_prime_end   = g;
  }

  // Loop through upscatter.
  for (int g_prime = g_prime_begin; g_prime < g_prime_end; g_prime++)
  {
    // Grab the g_prime flux moments.
    const Moments_Field_t &phi_g_prime = d_state->get_phi(g_prime);

    // Loop through cells and unknowns so we can calculate S * v
    // iteration-vector-flux moments at each unknown location.
    for (int cell = 0; cell < d_mesh->number_cells(); cell++)
    {
      // Grab the cross sections for this cell.
      const typename Materials::XS &xs = d_materials->get(matid[cell]);

      // Check if we have upscattering from g_prime to g.  If
      // not, then we continue to the next cell.  We could do all groups for
      // a cell, but that probably is way less cache-efficient.
      if (xs.group_upper_bound(g) < g_prime) continue;

      // Loop through spatial unknowns in the cell
      for (int n = 0, N = Equations::number_unknowns(); n < N; n++)
      {
        // Loop through moments
        for (int l = 0; l <= phi_g_prime.legendre_order(); l++)
        {

          // Get sigma_s for g<-g_prime, l
          double sig = xs.sigma_s(g, g_prime, l);

          // Get the m indices for this l
          const std::vector<int> &m_index = d_momentsindex->m_index(l);

          // Loop over m
          for (int i = 0; i < m_index.size(); i++)
          {
            int m = m_index[i];
            // Add the prime contribution.
            s(cell, n, l, m) += sig * phi_g_prime(cell, n, l, m);
          }
        }

      }
    } // end cell loop
  } // end upscatter loop
  return;
}

// Set within-group source, i.e. s_g = sum_g' [S_gg' phi_g'], where g' != g
template <class Equations, class D>
void SweepSource<Equations, D>::
build_within_group_source(int g, const Moments_Field_t &phi_g)
{

  // Make short reference to within group source
  Moments_Field_t &s = d_within_group_source;

  // Grab cell material id's.
  const typename Mesh<D>::Vec_Id &matid = d_mesh->get_cell_id(MATERIAL_ID);

  // Loop through cells and unknowns so we can calculate S * v
  // iteration-vector-flux moments at each unknown location.
  for (int cell = 0; cell < d_mesh->number_cells(); cell++)
  {

    // Grab the cross sections for this cell.
    const typename Materials::XS &xs = d_materials->get(matid[cell]);

    // Loop through spatial unknowns in the cell
    for (int n = 0, N = Equations::number_unknowns(); n < N; n++)
    {

      // Loop through moments
      for (int l = 0; l <= phi_g.legendre_order(); l++)
      {
        // Get sigma_s for g<-g,l (within-group scattering)
        double sig = xs.sigma_s(g, g, l);

        // Get the m indices for this l
        const std::vector<int> &m_index = d_momentsindex->m_index(l);

        // Loop over m
        for (int i = 0; i < m_index.size(); i++)
        {
          int m = m_index[i];
          s(cell, n, l, m) = sig * phi_g(cell, n, l, m);
        }
      }
    }
  }
  return;
}

// Combine everything into one source for one angle at each unknown.
template <class Equations, class D>
void SweepSource<Equations, D>::
build_source(int g, int o, int a)
{
  // Make short reference to sweep source
  SweepSourceField &s = d_source;

  // Grab the appropriate cardinal angle index for accessing M operator row.
  int angle = d_angularmesh->index(o, a);

  // Grab the corresponding Omega.
  const Omega &omega = (d_angularmesh->octant(o)).angles(a);

  /*
   * \note I'm using a straightforward loop to apply M for now.  We could
   *       (and later should) access the row of M directly along with
   *       the phi_g segments (directly).  These chunks of vectors could
   *       be dotted via BLAS.
   */

  // Set source to zero.
  std::fill(s.begin(), s.end(), 0.0);

  // Loop through all cells.
  for (int cell = 0; cell < d_mesh->number_cells(); cell++)
  {
    // Loop through spatial unknowns in the cell
    for (int n = 0, N = Equations::number_unknowns(); n < N; n++)
    {
      // Loop through moments (room for improvement)
      for (int l = 0; l <= d_momentsindex->legendre_order(); l++)
      {
        // Get m indices for this l.
        const std::vector<int> &m_index = d_momentsindex->m_index(l);

        // Loop over m
        for (int i = 0; i < m_index.size(); i++)
        {
          int m = m_index[i];
          // Add the scattering source contributions.
          s(cell, n) += (*d_m_operator)(angle, l, m)
              * (  d_within_group_source(cell, n, l, m)
                 + d_in_scatter_source(cell, n, l, m)   );
        }
      }

      // Add external source, if present.
      if (d_externalsource)
      {
        s(cell, n) += d_externalsource->q_e(cell, omega, g);
      }

      // Add fission source, if present.
      if (d_fissionsource)
      {
        s(cell, n) += d_fissionsource->q_f(cell, n, g);
      }

    } // end unknown loop
  } // end cell loop
  return;
}

// Instantiations.
template class SweepSource<DD_1D_Equations,_1D>;
template class SweepSource<DD_2D_Equations,_2D>;
template class SweepSource<DD_3D_Equations,_3D>;

} // end namespace slabtran

//---------------------------------------------------------------------------//
//              end of SweepSource.cc
//---------------------------------------------------------------------------//
